package com.ruoyi.common.utils;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.MessageManager;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import org.springframework.stereotype.Component;

import java.util.LinkedHashMap;
import java.util.Map;

@Component
public class ALiYunAISDKUtils {

    String apiKey = System.getenv("DASHSCOPE_API_KEY");
    String workspace = System.getenv("DASHCOPE_WORK_SPACE");

    public Map<String, String> customerWithOthers(String message) throws NoApiKeyException, InputRequiredException {

        if (apiKey.isEmpty() || workspace.isEmpty()){
            throw new NoApiKeyException();
        }
        Generation gen = new Generation();
        MessageManager msgManager = new MessageManager(10);
        Message systemMsg =
                Message.builder().role(Role.SYSTEM.getValue()).content("你是一个打车平台乐于助人的客服，并且你的回答简洁并且准确").build();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(message).build();
        msgManager.add(systemMsg);
        msgManager.add(userMsg);
        GenerationParam param =
                GenerationParam.builder().model(Generation.Models.QWEN_TURBO).messages(msgManager.get())
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .workspace(workspace)
                        .apiKey(apiKey)
                        .build();
        GenerationResult result = gen.call(param);
        String content = result.getOutput().getChoices().get(0).getMessage().getContent();
        Integer inputTokens = result.getUsage().getInputTokens();
        Integer outputTokens = result.getUsage().getOutputTokens();
        Integer totalTokens = result.getUsage().getTotalTokens();
        LinkedHashMap<String, String> map = new LinkedHashMap<>();

        // 返回的内容
        map.put("content", content);
        // 输入tokens
        map.put("inputTokens", inputTokens.toString());
        // 输出tokens
        map.put("outputTokens", outputTokens.toString());
        // 总tokens
        map.put("totalTokens", totalTokens.toString());


        return map;
    }

}
